{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "import torch\n",
    "import torch.nn as nn\n",
    "import torch.nn.functional as F\n",
    "from torch.utils.data import Subset\n",
    "from torchvision import datasets, transforms\n",
    "from complexPyTorch.complexLayers import ComplexBatchNorm2d, ComplexConv2d, ComplexLinear\n",
    "from complexPyTorch.complexLayers import ComplexDropout2d, NaiveComplexBatchNorm2d\n",
    "from complexPyTorch.complexLayers import ComplexBatchNorm1d\n",
    "from complexPyTorch.complexFunctions import complex_relu, complex_max_pool2d"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [],
   "source": [
    "batch_size = 64\n",
    "n_train = 1000\n",
    "n_test = 100\n",
    "trans = transforms.Compose([transforms.ToTensor(), transforms.Normalize((0.5,), (1.0,))])\n",
    "train_set = datasets.MNIST('../data', train=True, transform=trans, download=True)\n",
    "train_set = Subset(train_set, torch.arange(n_train))\n",
    "test_set = datasets.MNIST('../data', train=False, transform=trans, download=True)\n",
    "test_set = Subset(test_set, torch.arange(n_test))\n",
    "\n",
    "train_loader = torch.utils.data.DataLoader(train_set, batch_size= batch_size, shuffle=True)\n",
    "test_loader = torch.utils.data.DataLoader(test_set, batch_size= batch_size, shuffle=True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [],
   "source": [
    "class ComplexNet(nn.Module):\n",
    "    \n",
    "    def __init__(self):\n",
    "        super(ComplexNet, self).__init__()\n",
    "        self.conv1 = ComplexConv2d(1, 10, 5, 1)\n",
    "        self.bn2d  = ComplexBatchNorm2d(10, track_running_stats = False)\n",
    "        self.conv2 = ComplexConv2d(10, 20, 5, 1)\n",
    "        self.fc1 = ComplexLinear(4*4*20, 500)\n",
    "        self.dropout = ComplexDropout2d(p = 0.3)\n",
    "        self.bn1d  = ComplexBatchNorm1d(500, track_running_stats = False)\n",
    "        self.fc2 = ComplexLinear(500, 10)\n",
    "             \n",
    "    def forward(self,x):\n",
    "        x = self.conv1(x)\n",
    "        x = complex_relu(x)\n",
    "        x = complex_max_pool2d(x, 2, 2)\n",
    "        x = self.bn2d(x)\n",
    "        x = self.conv2(x)\n",
    "        x = complex_relu(x)\n",
    "        x = complex_max_pool2d(x, 2, 2)\n",
    "        x = x.view(-1,4*4*20)\n",
    "        x = self.fc1(x)\n",
    "        x = self.dropout(x)\n",
    "        x = complex_relu(x)\n",
    "        x = self.bn1d(x)\n",
    "        x = self.fc2(x)\n",
    "        x = x.abs()\n",
    "        x =  F.log_softmax(x, dim=1)\n",
    "        return x\n",
    "    \n",
    "device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')\n",
    "model = ComplexNet().to(device)\n",
    "optimizer = torch.optim.SGD(model.parameters(), lr=5e-3, momentum=0.9)\n",
    "\n",
    "def train(model, device, train_loader, optimizer, epoch):\n",
    "    model.train()\n",
    "    for batch_idx, (data, target) in enumerate(train_loader):\n",
    "        data, target =data.to(device).type(torch.complex64), target.to(device)\n",
    "        optimizer.zero_grad()\n",
    "        output = model(data)\n",
    "        loss = F.nll_loss(output, target)\n",
    "        loss.backward()\n",
    "        optimizer.step()\n",
    "        if batch_idx % 100 == 0:\n",
    "            print('Train\\t Epoch: {:3} [{:6}/{:6} ({:3.0f}%)]\\tLoss: {:.6f}'.format(\n",
    "                epoch,\n",
    "                batch_idx * len(data), \n",
    "                len(train_loader.dataset),\n",
    "                100. * batch_idx / len(train_loader), \n",
    "                loss.item())\n",
    "            )\n",
    "            \n",
    "def test(model, device, test_loader, optimizer, epoch):\n",
    "    model.eval()\n",
    "    for batch_idx, (data, target) in enumerate(train_loader):\n",
    "        data, target = data.to(device).type(torch.complex64), target.to(device)\n",
    "        output = model(data)\n",
    "        loss = F.nll_loss(output, target)\n",
    "        if batch_idx % 100 == 0:\n",
    "            print('Test\\t Epoch: {:3} [{:6}/{:6} ({:3.0f}%)]\\tLoss: {:.6f}'.format(\n",
    "                epoch,\n",
    "                batch_idx * len(data), \n",
    "                len(test_loader.dataset),\n",
    "                100. * batch_idx / len(test_loader), \n",
    "                loss.item())\n",
    "            )"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/opt/miniconda/envs/py38/lib/python3.8/site-packages/torch/autograd/__init__.py:130: UserWarning: Casting complex values to real discards the imaginary part (Triggered internally at  /opt/conda/conda-bld/pytorch_1607370124688/work/aten/src/ATen/native/Copy.cpp:162.)\n",
      "  Variable._execution_engine.run_backward(\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Train\t Epoch:   0 [     0/  1000 (  0%)]\tLoss: 2.575082\n",
      "Test\t Epoch:   0 [     0/  1000 (  0%)]\tLoss: 0.868310\n",
      "Train\t Epoch:   1 [     0/  1000 (  0%)]\tLoss: 0.803982\n",
      "Test\t Epoch:   1 [     0/  1000 (  0%)]\tLoss: 0.194764\n",
      "Train\t Epoch:   2 [     0/  1000 (  0%)]\tLoss: 0.366340\n",
      "Test\t Epoch:   2 [     0/  1000 (  0%)]\tLoss: 0.160019\n",
      "Train\t Epoch:   3 [     0/  1000 (  0%)]\tLoss: 0.411020\n",
      "Test\t Epoch:   3 [     0/  1000 (  0%)]\tLoss: 0.088335\n"
     ]
    }
   ],
   "source": [
    "# Run training on 4 epochs\n",
    "for epoch in range(4):\n",
    "    train(model, device, train_loader, optimizer, epoch)\n",
    "    test(model, device, test_loader, optimizer, epoch)"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.8.5"
  },
  "toc": {
   "base_numbering": 1,
   "nav_menu": {},
   "number_sections": true,
   "sideBar": true,
   "skip_h1_title": true,
   "title_cell": "Table of Contents",
   "title_sidebar": "Contents",
   "toc_cell": false,
   "toc_position": {},
   "toc_section_display": true,
   "toc_window_display": false
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}
